package org.ljm.basic.util;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStream;
import java.sql.SQLException;

import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.database.QueryDataSet;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.dataset.xml.FlatXmlProducer;
import org.dbunit.operation.DatabaseOperation;
import org.junit.AfterClass;
import org.junit.Assert;
import org.junit.BeforeClass;
import org.xml.sax.InputSource;

public class AbstractDbunitTestCase {

	public static IDatabaseConnection dbCon;
	
	private File tmpFile;
	
	@BeforeClass
	public static void init() throws DatabaseUnitException, SQLException{
		dbCon = new DatabaseConnection(DbUtil.getConn());
	}
	
	protected void backupAllTables() throws SQLException, DataSetException, IOException {
		IDataSet ds = dbCon.createDataSet();
		write2backupFile(ds);
	}

	private void write2backupFile(IDataSet ds) throws IOException,
			DataSetException {
		tmpFile = File.createTempFile("back", "xml");
		FlatXmlDataSet.write(ds, new FileWriter(tmpFile));
	}
	
	protected void backupCustomTables(String[] names) 
			throws DataSetException, IOException{
		QueryDataSet ds = new QueryDataSet(dbCon);
		for(String name : names){
			ds.addTable(name);
		}
		write2backupFile(ds);
	}
	
	protected void backupOneTable(String tname) throws DataSetException, IOException{
		backupCustomTables(new String[]{tname});
	}
	
	protected void resumeTables() throws FileNotFoundException, DatabaseUnitException, SQLException{
		IDataSet ds = new FlatXmlDataSet(new FlatXmlProducer(new InputSource(new FileInputStream(tmpFile))));
	    DatabaseOperation.CLEAN_INSERT.execute(dbCon, ds);
	}
	
	protected IDataSet createDataSet(String tname) throws DataSetException{
		InputStream is = getClass().getClassLoader()
				.getResourceAsStream(tname+".xml");
		Assert.assertNotNull("用于测试的数据文件不存在",is);
		return new FlatXmlDataSet(new FlatXmlProducer(new InputSource(is)));
	}
	
	@AfterClass
	public static void destory(){
		try {
			if(dbCon != null) dbCon.close();
		} catch (SQLException e) {
			e.printStackTrace();
		}
	}
}
